#!/usr/bin/env python3
# 版本: 0.5.0.20240803
# 作者: 陈振玺
# 功能: 进程终止工具，支持杀掉指定进程及其所有子进程

import os
import sys
import time
import signal
import psutil
import argparse
from typing import List, Set


class ProcessKiller:
    """进程终止器"""
    
    def __init__(self, force_kill: bool = False, wait_timeout: int = 5):
        """
        初始化进程终止器
        
        Args:
            force_kill: 是否强制杀掉进程（使用SIGKILL）
            wait_timeout: 等待进程正常退出的超时时间（秒）
        """
        self.force_kill = force_kill
        self.wait_timeout = wait_timeout
    
    def get_process_tree(self, pid: int) -> List[psutil.Process]:
        """
        获取进程树中的所有进程（包括父进程和所有子进程）
        
        Args:
            pid: 父进程PID
            
        Returns:
            进程列表，按照从子进程到父进程的顺序排列
        """
        try:
            parent = psutil.Process(pid)
            # 递归获取所有子进程
            children = parent.children(recursive=True)
            
            # 返回所有进程，子进程在前，父进程在后（确保先杀子进程）
            all_processes = children + [parent]
            
            print(f"📋 发现进程树: 父进程PID={pid}, 子进程数量={len(children)}")
            for i, proc in enumerate(all_processes):
                try:
                    print(f"  {i+1}. PID={proc.pid}, 名称={proc.name()}, 状态={proc.status()}")
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    print(f"  {i+1}. PID={proc.pid}, 名称=未知, 状态=无法访问")
            
            return all_processes
            
        except psutil.NoSuchProcess:
            print(f"❌ 进程 PID={pid} 不存在")
            return []
        except psutil.AccessDenied:
            print(f"❌ 无权限访问进程 PID={pid}")
            return []
        except Exception as e:
            print(f"❌ 获取进程树失败: {str(e)}")
            return []
    
    def kill_process(self, process: psutil.Process) -> bool:
        """
        终止单个进程
        
        Args:
            process: 要终止的进程对象
            
        Returns:
            是否成功终止进程
        """
        try:
            pid = process.pid
            name = process.name()
            
            # 检查进程是否还存在
            if not process.is_running():
                print(f"✅ 进程 PID={pid} ({name}) 已经不存在")
                return True
            
            if self.force_kill:
                # 强制杀掉进程
                print(f"🔥 强制终止进程 PID={pid} ({name})")
                process.kill()  # 发送SIGKILL信号
            else:
                # 优雅地终止进程
                print(f"🛑 优雅终止进程 PID={pid} ({name})")
                process.terminate()  # 发送SIGTERM信号
                
                # 等待进程退出
                try:
                    process.wait(timeout=self.wait_timeout)
                    print(f"✅ 进程 PID={pid} ({name}) 已正常退出")
                    return True
                except psutil.TimeoutExpired:
                    print(f"⏰ 进程 PID={pid} ({name}) 在{self.wait_timeout}秒内未退出，强制杀掉")
                    process.kill()  # 超时后强制杀掉
            
            # 再次等待确认进程已终止
            try:
                process.wait(timeout=3)
                print(f"✅ 进程 PID={pid} ({name}) 已被终止")
                return True
            except psutil.TimeoutExpired:
                print(f"❌ 进程 PID={pid} ({name}) 仍然存在")
                return False
                
        except psutil.NoSuchProcess:
            print(f"✅ 进程 PID={process.pid} 已经不存在")
            return True
        except psutil.AccessDenied:
            print(f"❌ 无权限终止进程 PID={process.pid}")
            return False
        except Exception as e:
            print(f"❌ 终止进程 PID={process.pid} 失败: {str(e)}")
            return False
    
    def kill_process_tree(self, pid: int) -> bool:
        """
        终止进程树（包括指定进程及其所有子进程）
        
        Args:
            pid: 要终止的父进程PID
            
        Returns:
            是否成功终止所有进程
        """
        print(f"🎯 开始终止进程树，根进程PID={pid}")
        
        # 获取进程树
        processes = self.get_process_tree(pid)
        if not processes:
            return False
        
        # 记录终止结果
        success_count = 0
        total_count = len(processes)
        
        print(f"\n🚀 开始终止 {total_count} 个进程...")
        
        # 逐个终止进程（从子进程开始）
        for i, process in enumerate(processes, 1):
            print(f"\n[{i}/{total_count}] 正在处理进程:")
            if self.kill_process(process):
                success_count += 1
        
        # 输出结果统计
        print(f"\n📊 终止结果统计:")
        print(f"  总进程数: {total_count}")
        print(f"  成功终止: {success_count}")
        print(f"  失败数量: {total_count - success_count}")
        
        if success_count == total_count:
            print(f"\n🎉 所有进程已成功终止！")
            return True
        else:
            print(f"\n⚠️ 部分进程终止失败，请检查权限或手动处理")
            return False
    
    def check_process_exists(self, pid: int) -> bool:
        """
        检查进程是否存在
        
        Args:
            pid: 进程PID
            
        Returns:
            进程是否存在
        """
        try:
            process = psutil.Process(pid)
            if process.is_running():
                print(f"📍 进程信息: PID={pid}, 名称={process.name()}, 状态={process.status()}")
                return True
            else:
                print(f"❌ 进程 PID={pid} 不存在或已退出")
                return False
        except psutil.NoSuchProcess:
            print(f"❌ 进程 PID={pid} 不存在")
            return False
        except psutil.AccessDenied:
            print(f"❌ 无权限访问进程 PID={pid}")
            return False
        except Exception as e:
            print(f"❌ 检查进程失败: {str(e)}")
            return False


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="进程终止工具，支持杀掉指定进程及其所有子进程",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  process-killer -p 1234                    # 优雅终止进程1234及其子进程
  process-killer -p 1234 -f                # 强制终止进程1234及其子进程
  process-killer -p 1234 -t 10             # 设置10秒超时时间
  process-killer -p 1234 -c                # 仅检查进程是否存在
        """
    )
    
    parser.add_argument(
        "-p", "--pid",
        type=int,
        required=True,
        help="要终止的进程PID"
    )
    
    parser.add_argument(
        "-f", "--force",
        action="store_true",
        help="强制杀掉进程（使用SIGKILL，默认使用SIGTERM）"
    )
    
    parser.add_argument(
        "-t", "--timeout",
        type=int,
        default=5,
        help="等待进程正常退出的超时时间（秒，默认: 5）"
    )
    
    parser.add_argument(
        "-c", "--check",
        action="store_true",
        help="仅检查进程是否存在，不执行终止操作"
    )
    
    parser.add_argument(
        "-y", "--yes",
        action="store_true",
        help="跳过确认提示，直接执行"
    )
    
    args = parser.parse_args()
    
    print(f"🔧 进程终止工具启动")
    print(f"🎯 目标进程PID: {args.pid}")
    print(f"⚡ 终止模式: {'强制终止(SIGKILL)' if args.force else '优雅终止(SIGTERM)'}")
    print(f"⏰ 超时时间: {args.timeout}秒")
    print("-" * 50)
    
    # 创建进程终止器
    killer = ProcessKiller(force_kill=args.force, wait_timeout=args.timeout)
    
    # 检查进程是否存在
    if not killer.check_process_exists(args.pid):
        sys.exit(1)
    
    # 如果只是检查模式，直接退出
    if args.check:
        print(f"\n✅ 进程 PID={args.pid} 存在")
        sys.exit(0)
    
    # 获取进程树信息
    processes = killer.get_process_tree(args.pid)
    if not processes:
        sys.exit(1)
    
    # 确认操作
    if not args.yes:
        print(f"\n⚠️ 警告: 即将终止 {len(processes)} 个进程")
        print(f"终止模式: {'强制终止(SIGKILL)' if args.force else '优雅终止(SIGTERM)'}")
        
        try:
            confirm = input("\n确认执行吗？(y/N): ").strip().lower()
            if confirm not in ['y', 'yes']:
                print("❌ 操作已取消")
                sys.exit(0)
        except KeyboardInterrupt:
            print("\n❌ 操作已取消")
            sys.exit(0)
    
    # 执行终止操作
    success = killer.kill_process_tree(args.pid)
    
    if success:
        print(f"\n🎉 进程树终止完成！")
        sys.exit(0)
    else:
        print(f"\n❌ 进程树终止失败！")
        sys.exit(1)


if __name__ == "__main__":
    main()